

#include <stdio.h>
#include <stdlib.h>

#include <assert.h>
#include <math.h>

#include <mpi.h>

#include "utils.h"
#include "mv.h"

//#define DEBUG

void mv(base_t **A, int nrows, int ncols, int nrows_a_loc, int ncols_a_loc,
        base_t *x, int nrows_x_loc,
        base_t *b, int ncols_b_loc)
{
    int comm_rank, comm_size, ret;

    // query communicator parameters
    MPI_Comm_rank(MPI_COMM_WORLD, &comm_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &comm_size);

#ifdef DEBUG
    printf("p_%d: (nrows, ncols, nrows_a_loc, ncols_a_loc, nrows_x_loc, ncols_b_loc): (%d, %d, %d, %d, %d, %d)\n", comm_rank, nrows, ncols, nrows_a_loc, ncols_a_loc, nrows_x_loc, ncols_b_loc);
    for (int i = 0; i < nrows_x_loc; i++) {
        printf("p_%d: xloc[%d]: %.2f\n", comm_rank, i, x[i]);
    }
    for (int j = 0; j < nrows_a_loc; j++) {
        printf("p_%d: a[%d]: ", comm_rank, j);
        for (int i = 0; i < ncols_a_loc; i++) {
            printf("%.2f ", A[j][i]);
        }
        printf("\n");
    }
#endif

    // we need to use the amount of matrix rows (total amount of
    // b vector rows!)
    base_t b_full[nrows];
    for (int i = 0; i < nrows; i++) {
        b_full[i] = 0;
    }

    for (int j = 0; j < nrows_a_loc; j++) {
        for (int i = 0; i < ncols_a_loc; i++) {
            b_full[j] += A[j][i] * x[i];
        }
    }

#ifdef DEBUG
    for (int i = 0; i < nrows; i++) {
        printf("p_%d: b[%d]: %.2f\n", comm_rank, i, b_full[i]);
    }
#endif

    // synchronize receive counts for all processes before using
    // collective reduce scatter operation
    int recv_counts[comm_size];
    ret = MPI_Allgather(
        // send buffer, count and type
        &ncols_b_loc, 1, MPI_INT,
        // recv buffer, count and type
        recv_counts, 1, MPI_INT,
        // communicator
        MPI_COMM_WORLD
    );

#ifdef DEBUG
    printf("p_%d: allgather returned %d\n", comm_rank, ret);
    printf("p_%d: ", comm_rank);
    for (int i = 0; i < comm_size; i++) {
        printf("%d ", recv_counts[i]);
    }
    printf("\n");
#endif

    // use reduce scatter to reduce all b_full vectors together and then
    // distribute the b_full vector with varying counts to all processes
    ret = MPI_Reduce_scatter(
        // send buffer
        b_full,
        // recv buffer, count and type
        b, recv_counts, MPI_DOUBLE,
        // reduction operation
        MPI_SUM,
        // communicator
        MPI_COMM_WORLD
    );
}
